package it.uniroma2.dtk.drp;

import it.uniroma2.dtk.dt.GenericDT;
import it.uniroma2.util.math.ArrayMath;
import it.uniroma2.util.tree.Tree;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;

public class LearningMatricesBuilder {
	
	public static void main(String [] argv) throws Exception {
		GenericDT fft = new GenericDT(0, Integer.parseInt(argv[1]), true, false, GenericDT.Types.CONVOLUTION);
		fft.setLambda(0);
		LearningMatricesBuilder lmb = new LearningMatricesBuilder(fft);
		File treeFile = new File(argv[0]);
		lmb.processTreeFile(treeFile);
	}

	private GenericDT ts = null;

	public LearningMatricesBuilder(GenericDT dtfGenerator) {
		ts = dtfGenerator;
	}

	public void processTreeFile(File treeFile) throws Exception {
		BufferedWriter bops = new BufferedWriter(new FileWriter(new File(treeFile.getParentFile(), treeFile.getName() + "_BOP.svm")));
		BufferedWriter bops2 = new BufferedWriter(new FileWriter(new File(treeFile.getParentFile(), treeFile.getName() + "_BOP+BI.svm")));
		BufferedWriter bops3 = new BufferedWriter(new FileWriter(new File(treeFile.getParentFile(), treeFile.getName() + "_BOP+TRI.svm")));
//		BufferedWriter dtfs = new BufferedWriter(new FileWriter(new File(treeFile.getParentFile(), treeFile.getName() + "_DTF.svm")));
		BufferedReader in = new BufferedReader(new FileReader(treeFile));

		String line = in.readLine();
		int count = 0;
		while (line != null) {
			System.out.print(".");
			if (++count % 100 == 0)
				System.out.println(count);
			String parsedSentence = line.substring(line.indexOf("|BT|")+4, line.indexOf("|ET|")).trim();
			Tree x = Tree.fromPennTree(parsedSentence);
			polishPennTree(x);

//			double[] dtf = ts.cds(x);
			ArrayList<double[]> dpos = dpos(x);

			writeVector(bops, dpos.get(0));
			writeVector(bops2, dpos.get(1));
			writeVector(bops3, dpos.get(2));
//			writeVector(dtfs, dtf);
			line = in.readLine();
		}
		in.close();
		bops.close();
		bops2.close();
		bops3.close();
//		dtfs.close();
	}

	private void writeVector(BufferedWriter out, double[] dpos) throws IOException {
		out.write("-1\t" + (1) + ":" + dpos[0]);
		for (int i=1; i<dpos.length; i++) {
			out.write( " " + (i+1) + ":" + dpos[i]);
		}
		out.write("\n");
	}

	private ArrayList<double[]> dpos(Tree tree) throws Exception {
		ArrayList<Tree> preterminals = new ArrayList<Tree>();
		getPreterminals(tree, preterminals);
		ArrayList<double[]> result = new ArrayList<double[]>(3);
		double[] out1 = null;
		double[] out2 = null;
		double[] out3 = null;
		for (int i = preterminals.size(); i > 0; i--) {
			if (out1 == null) {
				out1 = ts.getLabelVector(preterminals.get(i-1));
				out2 = Arrays.copyOf(out1, out1.length);
				out3 = Arrays.copyOf(out1, out1.length);
			}
			else {
				//BagOfPOS:unigrams+bigrams
				// UNIGRAMMA
				double[] uni = ArrayMath.sum(out1, ts.getLabelVector(preterminals.get(i-1)));
				out1 = ArrayMath.sum(out1 , uni);
				out2 = ArrayMath.sum(out2 , uni);
				out3 = ArrayMath.sum(out3 , uni);
				// BIGRAMMA
				double[] bi = ts.op(ts.getLabelVector(preterminals.get(i-1)),ts.getLabelVector(preterminals.get(i)));		
				out2 = ArrayMath.sum( out2 , bi );
				out3 = ArrayMath.sum( out3 , bi );
				// TRIGRAMMA
				if (i-2 >= 0) {
					double[] tri = ts.op(ts.getLabelVector(preterminals.get(i-2)),bi);
					out3 = ArrayMath.sum( out1 , tri );
				}
			}
		}
		result.add(out1);
		result.add(out2);
		result.add(out3);
		return result;
	}
	
	private void getPreterminals(Tree tree, ArrayList<Tree> preterminals) {
		if (tree.getChildren().size() == 1 && tree.getChildren().get(0).isTerminal())
			preterminals.add(tree);
		else
			for (Tree child : tree.getChildren())
				getPreterminals(child, preterminals);
	}

	private Tree polishPennTree(Tree tree) {
		Tree out = null;
		if (tree.getChildren().size() == 1 && tree.getChildren().get(0).isTerminal()) {
			if (!tree.getRootLabel().trim().contains("*")) 
				out = tree; 
		} else {
			ArrayList<Tree> newChildren = new ArrayList<Tree>();
			for (Tree o : tree.getChildren()) {
				Tree newChild = polishPennTree(o);
				if (newChild!=null) 
					newChildren.add(newChild);
			}
			if (newChildren.size() != 0) {
				tree.getChildren().clear();
				tree.getChildren().addAll(newChildren);
				out = tree;
			}
		}
		if (out != null) {
			if (tree.getRootLabel().contains("-")) tree.setRootLabel(tree.getRootLabel().substring(0,tree.getRootLabel().indexOf('-')) ); 
		}
		return out;
	}

}
